# -*- coding: utf-8 -*-

# File Name : transform
# Description :
# Author : marje
# date : 2022/6/5
# Change Activity : 2022/6/5
import numpy
import numpy as np
import torch
from torchvision.transforms.functional import to_tensor, normalize

from ml.dl.misc.img_augment.resize_img import resize_img


class ImageToTensor:
    """
    change ndarray type to tensor.
    """

    def __init__(self, input_key='img', output_key='img'):
        self.input_key = input_key
        self.output_key = output_key

    @staticmethod
    def handle(img):
        """


        Args:
            img:one image with shape (w,h,c) or multiple images with shape (b,w,h,c)

        Returns:
            Tensor: tensor type object with shape (b,c,w,h)
        """
        tmp_img = img.copy()
        return to_tensor(tmp_img)

    def __call__(self, data):
        """

        Args:
            data:

        Returns:

        """
        img = data[self.input_key]
        res = ImageToTensor.handle(img)
        data[self.output_key] = res
        return data


class TensorToDevice:
    """
        change ndarray type to tensor.
    """

    def __init__(self, input_key='img', output_key='img', device_name='cuda:0'):
        """

        Args:
            input_key:
            output_key:
            device_name:
        """
        self.input_key = input_key
        self.output_key = output_key
        self.device_name = device_name

    @staticmethod
    def handle(x: torch.Tensor, device_name):
        """
        input tensor to device

        Args:
            x: input tensor
            device_name: device name. eg: cpu,cuda:0
        Returns:
            Tensor
        """

        device = torch.device(device_name)
        y = x.to(device)
        return y

    def __call__(self, data):
        """

        Args:
            data:

        Returns:

        """
        img = data[self.input_key]
        res = TensorToDevice.handle(img, self.device_name).contiguous()
        data[self.output_key] = res
        return data


class ToDataType:
    def __init__(self, input_key='img', output_key='img', dtype=torch.float32):
        """

        Args:
            input_key:
            output_key:
            dtype: ndarray or tensor type.
        """
        self.input_key = input_key
        self.output_key = output_key
        self.dtype = dtype

    @staticmethod
    def handle(x, dtype):
        """
        input tensor to device

        Args:
            x: input
            dtype: destination data type. eg: np.float32, torch.float32
        Returns:
            Tensor
        """
        if isinstance(x, np.ndarray):
            y = x.astype(dtype=dtype)
            y = np.ascontiguousarray(y)
        else:
            y = x.to(dtype).contiguous()
        return y

    def __call__(self, data):
        """

        Args:
            data:

        Returns:

        """
        img = data[self.input_key]
        res = ToDataType.handle(img, self.dtype)
        data[self.output_key] = res
        return data


class ResizeFixHeightTextarea:
    def __init__(self, textarea_key='img', output_key='img', fixed_height=32, max_width=320):
        """
        Resize textarea.
        Args:
            textarea_key:
            output_key:
            fixed_height:
            max_width:
        """
        self.textarea_key = textarea_key
        self.output_key = output_key
        self.fixed_height = fixed_height
        self.max_width = max_width

    @staticmethod
    def handle(img, fixed_height=32, max_width=320):
        ori_h, ori_w = img.shape[:2]
        ratio = ori_h / fixed_height
        new_width = min(int(ori_w / ratio), max_width)
        if new_width % 2 == 1:  # if width is odd, turn it into an even number
            new_width = new_width + 1
        return resize_img(img, (new_width, fixed_height))


class Normalize(object):
    def __init__(self, mean, std, input_key='img', output_key='img'):
        self.mean = mean
        self.std = std
        self.input_key = input_key
        self.output_key = output_key

    def __call__(self, data):
        img = data[self.input_key]
        data[self.output_key] = normalize(img, mean=self.mean, std=self.std)
        return data
